﻿#ifndef ___PCLIB_S_WEBSOCKET___
#define ___PCLIB_S_WEBSOCKET___


#include "../m_core/m_core.hpp"




/**
*@brief    websocket协议解析类
*       +-+-+-+-+-------+-+-------------+-------------------------------+
*        0                   1                   2                   3
*        0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
*       +-+-+-+-+-------+-+-------------+-------------------------------+
*       |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
*       |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
*       |N|V|V|V|       |S|             |   (if payload len==126/127)   |
*       | |1|2|3|       |K|             |                               |
*       +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
*       |     Extended payload length continued, if payload len == 127  |
*       + - - - - - - - - - - - - - - - +-------------------------------+
*       |                     Payload Data continued ...                |
*       +---------------------------------------------------------------+
*/
class CPCWebSocketUnPack : CPCNoCopyable
{
public:
    /**
    *@brief        帧头1字节
    */
    enum PCWSFrameType
    {
        PCWS_TEXT_FRAME = 0x81,    //文本数据帧
        PCWS_BIN_FRAME = 0x82,    //二进制数据帧
        PCWS_CLOSE_FRAME = 0x88,    //关闭数据帧
        PCWS_PING_FRAME = 0x89,    //心跳帧PING
        PCWS_PONG_FRAME = 0x8A     //心跳帧PONG
    };

    /**
    *@brief        构造方法
    *@param        bClient [in]    是否为websocket客户端，否则为服务端，客户端和服务端区别如下：
    *                            1.握手时客户端发起认证，服务端要根据客户端发起的认证码计算认证结果(公开算法)
    *                            2.客户端数据必须经过mask，服务端数据不需要
    */
    explicit CPCWebSocketUnPack(bool bClient)
        :m_bClient(bClient)
    {}

    /**
    *@brief     客户端生成Sec-WebSocket-Key
    *@return    结果
    */
    static std::string GenerateSecWebSocketKey()
    {
        std::string binkey = CPCRandom::PCGenString(16, "");
        return CPCStrUtil::PCBase64Encode((const unsigned char*)binkey.data(), binkey.length());
    }

    /**
    *@brief     服务端根据请求的Sec-WebSocket-Key生成Sec-WebSocket-Accept数据
    *@param     secWebSocketKey [in]    客户端给的密钥
    *@return    回应的校验值
    */
    static std::string GetValidateCode(const std::string& secWebSocketKey)
    {
        // 使用请求传过来的KEY+协议字符串，先用SHA1加密然后使用base64编码算出一个应答的KEY
        std::string serverKey = secWebSocketKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
        // SHA1 & Base64
        std::string digest = CPCSHA1::GetBytesSHA1(serverKey.data(),serverKey.length());
        return CPCStrUtil::PCBase64Encode((unsigned char*)digest.data(), PCSHA1_LEN);
    }

    /**
    *@brief     打包心跳帧
    *@param     bPing [in]    true为PING，false为PONG
    *@param     data  [in]    心跳帧附带的数据，一般使用时不应该附带数据（但是协议是支持的）
    *@return    打包结果数据。
    *@note      1.任意一方可主动发送PING(可附带数据)，收到PING必须回应相同的PONG（包括数据）
    *           2.没收到PING而主动发送PONG被认为是单向心跳，单向心跳不应该被回应
    *           3.如果收到PING没及时回应时又收到了第二个PING，可直接回应第二个PING，忽略第一个
    *           总结：收到PING必须回应相同数据的PONG，收到PONG什么都不要做
    */
    std::string PackHeartBeatFrame(bool bPing, const std::string& data = std::string())
    {
        if (!bPing)
        {
            return PackFrame(data, PCWS_PONG_FRAME);
        }
        else
        {
            return PackFrame(data, PCWS_PING_FRAME);
        }
    }

    /**
    *@brief     打包断开帧
    *@param     code   [in]    断开错误码，参考websocket标准错误码，如1000代表正常关闭
    *@param     reason [in]    断开原因（可选，必须传入UTF8编码）
    *@return    打包结果数据。
    */
    std::string PackCloseFrame(unsigned short code = 1000, const char* reason = NULL)
    {
        char szCode[2] = { 0 };
        code = htons(code);
        memcpy(szCode, &code, 2);
        std::string data(szCode, 2);
        if (reason != NULL &&  reason[0] != 0)
        {
            data += reason;
        }
        return PackFrame(data, PCWS_CLOSE_FRAME);
    }

    /**
    *@brief     打包帧
    *@param     data [in]    用户数据，不能超过4GB长度
    *@param     type [in]    帧类型。注意：文本帧（包括错误原因数据）需要使用UTF-8编码
    *@return    帧数据。
    */
    std::string PackFrame(const std::string& data, PCWSFrameType type = PCWS_TEXT_FRAME)
    {
        std::string frame;
        frame.reserve(data.length() + 16);
        frame.push_back((char)type);
        if (data.length() <= 125)
        {
            char secondbyte = (char)data.length();
            if (m_bClient) secondbyte = secondbyte | 0x80;
            frame.push_back(secondbyte);
        }
        else if (data.length() <= 65535)
        {
            char secondbyte = 126;
            if (m_bClient) secondbyte = secondbyte | 0x80;
            frame.push_back(secondbyte);
            frame.push_back((char)((data.length() >> 8) & 0xFF));
            frame.push_back((char)(data.length() & 0xFF));
        }
        else
        {
            char secondbyte = 127;
            if (m_bClient) secondbyte = secondbyte | 0x80;
            frame.push_back(secondbyte);

            uint64_t datalen = data.length();
            uint64_t size64 = PCHTON64(datalen);
            frame.append((const char*)(&size64), sizeof(size64));
        }
        //mask
        if (m_bClient)
        {
            std::string maskingkey = CPCRandom::PCGenString(4, "");
            frame.append(maskingkey);

            size_t datapos = frame.length();
            frame.resize(datapos + data.length());
            for (unsigned int i = 0; i < data.length(); i++)
            {
                int j = i % 4;
                frame[datapos + i] = data[i] ^ maskingkey[j];
            }
        }
        else
        {
            frame.append(data);
        }
        return frame;
    }

    /**
    *@brief     从流中读取并解包一帧，自动处理心跳和关闭帧
    *@param     socketptr  [in]    接收数据的socket
    *@param     nWaitMs    [in]    接收数据超时时间
    *@return    是否成功，成功则一帧数据保存在StrResult1，其他信息在成员变量中
    */
    CPCResult<std::string> UnPackFrame(CPCSocketHandle* socketptr, int nWaitMs)
    {
        CPCResult<std::string> result;
        std::string frame;
        while (1)
        {
            //收包头
            CPCBuffer headbuf(2);
            auto result2 = socketptr->XRecv(nWaitMs, headbuf, true);
            if (!result2)
            {
                return result.SetFail("XRecv fail:%s",result2.ErrDesc());
            }
            m_fin = (unsigned char)headbuf[0] >> 7;
            m_opcode = (unsigned char)headbuf[0] & 0x0f;
            m_mask = (unsigned char)headbuf[1] >> 7;
            m_payloadLen = (unsigned char)headbuf[1] & 0x7f;

            //处理合法性
            if (m_fin != 0 && m_fin != 1)
            {
                return result.SetFail("m_fin = %02x invalid", (int)m_fin);
            }
            if (!m_bClient && m_mask != 1)
            {
                //客户端过来的数据必须mask
                return result.SetFail("from client data must masked");
            }

            //更新m_payloadLen
            if (m_payloadLen == 126)
            {
                //收包长度2
                CPCBuffer lenbuf(2);
                auto result2 = socketptr->XRecv(nWaitMs, lenbuf, true);
                if (!result2)
                {
                    return result.SetFail("XRecv fail:%s",result2.ErrDesc());
                }
                uint16_t length = 0;
                memcpy(&length, lenbuf.data(), 2);
                m_payloadLen = ntohs(length);
            }
            else if (m_payloadLen == 127)
            {
                //收包长度8
                CPCBuffer lenbuf(8);
                auto result2 = socketptr->XRecv(nWaitMs, lenbuf, true);
                if (!result2)
                {
                    return result.SetFail("XRecv fail:%s",result2.ErrDesc());
                }
                uint64_t length = 0;
                memcpy(&length, lenbuf.data(), 8);
                m_payloadLen = PCNTOH64(length);

                //64位长度超过4GB了，32位程序不支持收这么长的数据
                if (m_payloadLen > PC_SIZE_T_MAX)
                {
                    return result.SetFail("m_payloadLen is too long, unsupported.");
                }
            }

            //收payload
            CPCBuffer payloadbuf((size_t)m_payloadLen + ((m_mask == 1) ? 4 : 0));
            if(payloadbuf.capacity() > 0)
            {
                //有payload才处理
                result2 = socketptr->XRecv(nWaitMs, payloadbuf, true);
                if (!result2)
                {
                    return result.SetFail("XRecv fail:%s",result2.ErrDesc());
                }
            }
            const char* payloaddata = payloadbuf.data();
            if(payloaddata == NULL)
                payloaddata = "";

            //根据类型处理
            if (m_opcode == 0x09)
            {
                //处理心跳PING
                std::string pingdata;
                AppendDecryptPayLoad(payloaddata, pingdata);

                //回复PONG，忽略send结果
                std::string pongframe = this->PackHeartBeatFrame(false, pingdata);
                socketptr->XSend(pongframe.data(), pongframe.length(), nWaitMs);
                continue;
            }
            else if (m_opcode == 0x0a)
            {
                //处理心跳PONG，不回复
                continue;
            }
            else if (m_opcode == 0x00 || m_opcode == 0x01 || m_opcode == 0x02)
            {
                //处理payload
                if (m_opcode != 0x00 && (!frame.empty()))
                {
                    //走到这里，说明在两个不完整分片之间插入了另外一个数据帧（或片，不是控制帧，协议规定了允许控制帧在分片之间）
                    //按照协议理解，不应该出现这种情况，否则无法区分两个帧都分片但是交错到达的情况
                    //这里的做法是简单地直接清理掉未完成的分片。
                    //TODO：可能丢失不完整的分片，要不要给用户一个处理协议意外但收到部分分片数据的回调？
                    //TODO：或者直接报错更严谨一些？
                    frame.clear();
                }
                AppendDecryptPayLoad(payloaddata, frame);
            }
            else if (m_opcode == 0x08)
            {
                //处理关闭帧
                std::string closedata;
                AppendDecryptPayLoad(payloaddata, closedata);

                unsigned short errcode = -1;
                std::string errdesc;
                if (closedata.length() >= 2)
                {
                    //可选的错误码
                    memcpy(&errcode, closedata.data(), 2);
                    errcode = ntohs(errcode);
                    char szErrDesc[64] = { 0 };
                    snprintf(szErrDesc, sizeof(szErrDesc) - 1, "remote closed.code=%hu,reason=", errcode);
                    errdesc = szErrDesc;

                    //可选的错误原因
                    if (closedata.length() > 2)
                    {
                        errdesc.append(closedata.data() + 2, closedata.length() - 2);
                    }
                }
                return result.SetFail(errcode,"%s",errdesc.c_str());
            }
            else
            {
                //不支持的帧类型
                return result.SetFail("m_opcode = %02x unsupported.", (int)m_opcode);
            }

            if (m_fin == 0)
            {
                //分片了，等待后面的片
                continue;
            }
            else
            {
                //最后一片或者只有一片
                break;
            }
        }
        return result.SetSucc(std::move(frame));
    }

protected:
    /**
    *@brief     解码payload并附加到输入的payloadbuf
    *@param     payloaddata  [in]     输入的数据
    *@param     payloadbuf   [out]    输出的缓冲区
    */
    inline void AppendDecryptPayLoad(const char* payloaddata, std::string& payloadbuf)
    {
        if (m_mask == 1)
        {
            //解密客户端数据
            unsigned char masking_key[4] = { 0, 0, 0, 0 };
            memcpy(masking_key, payloaddata, 4);

            size_t datapos = payloadbuf.length();
            payloadbuf.resize(datapos + (size_t)m_payloadLen);
            for (size_t i = 0; i < m_payloadLen; i++)
            {
                int j = i % 4;
                payloadbuf[datapos + i] = payloaddata[i + 4] ^ masking_key[j];
            }
        }
        else
        {
            payloadbuf.append(payloaddata, (size_t)m_payloadLen);
        }
    }

public:
    //解码器是否为websocket客户端进行解码
    bool m_bClient;

    //接收的包中的额外信息
    unsigned char m_fin;    //0后续还有数据 否则结束 
    unsigned char m_opcode; //0x0表示附加数据帧 0x1表示文本数据帧 0x2表示二进制数据帧 0x3-7暂时无定义，为以后的非控制帧保留 0x8表示连接关闭 0x9表示ping 0xA表示pong 0xB-F暂时无定义，为以后的控制帧保留
    unsigned char m_mask;   //用于标识PayloadData是否经过掩码处理，客户端发出的数据帧需要进行掩码处理为1
    uint64_t m_payloadLen;  //附带的数据长度，按7位，7+16位，7+64位递推
};




/**
*@brief        WebSocket请求类
*/
class CPCWebSocketClient : public CPCNoCopyable
{
public:
    explicit CPCWebSocketClient()
        :m_wsUnPack(true)
        , m_hSocketPtr(std::make_shared<CPCSocketHandle>(true, &CPCPoller::obj()))
    {
    }
    ~CPCWebSocketClient()
    {
        if (m_hSocketPtr)
            m_hSocketPtr->XClose();
    }

    /**
    *@brief     握手ws/wss请求
    *@param     url         [in]    请求url
    *@param     nTimeOut    [in]    超时毫秒。如果nTimeoutMs < 0 代表永不超时
    *@return    是否成功。 成功后用下面的取结果函数取结果
    */
    CPCResult<> WSHandeShake(const char* url, int nTimeOut)
    {
        CPCHttpUrl httpurl;
        auto result = httpurl.Parse(url);
        if (!result)
        {
            return result;
        }
        bool bssl = false;
        if (strcasecmp(httpurl.m_Schema.c_str(), "wss") == 0)
        {
            bssl = true;
        }
        CPCSockAddr addr;
        result = addr.GetHostByName(httpurl.m_Host.c_str(), httpurl.m_Port);
        if (!result)
        {
            return result;
        }
        auto result2 = m_hSocketPtr->XConnect(addr, nTimeOut, bssl);
        if (!result2)
        {
            return result.SetFail("XConnect fail:%s",result2.ErrDesc());
        }

        //客户端要生成随机密钥和验证结果
        std::string wsreqkey = CPCWebSocketUnPack::GenerateSecWebSocketKey();
        std::string wsreqacc = CPCWebSocketUnPack::GetValidateCode(wsreqkey);

        //组织数据并发送
        CPCHttpHeader header;
        header.SetRequestLine("GET", httpurl);
        header.SetHost(httpurl);
        header.SetHeader("Sec-WebSocket-Key", wsreqkey.c_str());
        header.SetHeader("Connection", "Upgrade");
        header.SetHeader("Upgrade", "websocket");
        header.SetHeader("Sec-WebSocket-Version", "13");
        header.SetHeader("Origin", "null");
        header.SetHeader("Cache-Control", "no-cache");
        std::string strheader = header.ToString(true, true);
        result = m_hSocketPtr->XSend(strheader.data(), strheader.length(), nTimeOut);
        if (!result)
        {
            m_hSocketPtr->XClose();
            return result;
        }
        CPCRestHttpUnPack httpUnPack;
        httpUnPack.UnPackInit(false);
        while (1)
        {
            CPCBuffer package(PC_MAXSTACK);
            result = m_hSocketPtr->XRecv(nTimeOut, package);
            if (!result)
            {
                m_hSocketPtr->XClose();
                return result;
            }
            httpUnPack.UnPackUpdate(package.data(), package.length());
            if(!httpUnPack.m_LastResult)
            {
                //出错
                m_hSocketPtr->XClose();
                return result.SetFail("%s",httpUnPack.m_LastResult.ErrDesc());
            }
            else if (0 == httpUnPack.m_LastResult.Get())
            {
                //继续处理
                continue;
            }
            else
            {
                //完成，验证是否为升级到websocket
                if (httpUnPack.m_header.m_resp_httpcode != 101)
                {
                    m_hSocketPtr->XClose();
                    return result.SetFail(httpUnPack.m_header.m_resp_httpcode,"upgrade to websocket not return 101(101 is upgrade succ, other is fail).");
                }
                //验证是否收到key accept
                auto wskeyacc = httpUnPack.m_header.m_header.find("Sec-WebSocket-Accept");
                if (wskeyacc == httpUnPack.m_header.m_header.end())
                {
                    //出错，服务端没有回应密钥
                    m_hSocketPtr->XClose();
                    return result.SetFail("server not response Sec-WebSocket-Accept .");
                }
                if ((*wskeyacc).second != wsreqacc)
                {
                    //出错，服务端回应密钥不一致
                    m_hSocketPtr->XClose();
                    return result.SetFail("server  response Sec-WebSocket-Accept not match.");
                }
                break;
            }
        }
        return result;
    }

    /**
    *@brief     发送用户数据，发送一帧
    *@param     data        [in]    要发送的用户数据
    *@param     nTimeoutMs  [in]    超时时间毫秒，<0代表永不超时
    *@param     type        [in]    帧类型，默认为文本帧
    *@return    是否成功
    */
    CPCResult<> SendText(const std::string& data, int nTimeoutMs = -1, CPCWebSocketUnPack::PCWSFrameType type = CPCWebSocketUnPack::PCWS_TEXT_FRAME)
    {
        std::string frame = m_wsUnPack.PackFrame(data, type);
        return m_hSocketPtr->XSend(frame.data(), frame.length(), nTimeoutMs);
    }

    /**
    *@brief     接收用户一帧数据帧（心跳帧会自动处理）
    *@param     nTimeOut      [in]    超时毫秒。如果nTimeoutMs < 0 代表永不超时
    *@return    是否成功(成功后帧数据：std::string，帧类型：可读取m_wsUnpack的成员变量进一步判断是文本还是二进制帧)
    */
    CPCResult<std::string> RecvText(int nTimeOut)
    {
        return m_wsUnPack.UnPackFrame(m_hSocketPtr.get(), nTimeOut);
    }

public:
    CPCWebSocketUnPack m_wsUnPack;
    CPCSocketHandlePtr m_hSocketPtr;
};

/**
*@brief      CWebSocket连接类
*            websoket有两次过程：1.http协议握手过程；2.握手后的数据帧收发过程
*            【本对象只能被智能指针管理】
*/
class CPCWebSocketConnection : public CPCSocketHandle
{
public:
    explicit CPCWebSocketConnection(int nWaitMs = 100000)
        :CPCSocketHandle(false, &CPCPoller::obj())
        , m_nWaitMs(nWaitMs)
        , m_wsUnPack(false){}

    /**
    *@brief     http协议握手完成后会调用这个函数
    *@return    返回成功后会继续往下处理，返回失败会直接中断连接
    */
    virtual CPCResult<> OnHandeShakeOver() { return CPCResult<>(); }

    /**
    *@brief     接收到用户数据帧的处理（心跳帧已经自动处理了），这个函数必须要实现
    *@param     data    [in]    收到的用户数据（可读取m_wsUnpack的成员变量进一步判断是文本还是二进制帧）。
    *@return    返回成功后会继续往下处理，返回失败会直接中断连接
    */
    virtual CPCResult<> OnRecvText(const std::string& data) { return CPCResult<>(); }

    /**
    *@brief     发送用户数据
    *@param     data        [in]    要发送的用户数据
    *@param     nTimeoutMs  [in]    超时时间毫秒，<0代表永不超时
    *@param     type        [in]    帧类型，默认为文本帧
    *@return    是否成功
    */
    CPCResult<> SendText(const std::string& data, int nTimeoutMs = -1, CPCWebSocketUnPack::PCWSFrameType type = CPCWebSocketUnPack::PCWS_TEXT_FRAME)
    {
        std::string frame = m_wsUnPack.PackFrame(data, type);
        return XSend(frame.data(), frame.length(), nTimeoutMs);
    }

    /**
    *@brief     线程池中处理数据，你可以修改这部分代码，定义自己的处理逻辑
    */
    void OnService() override
    {
        //处理握手过程
        auto result = HttpHandShake();
        if (!result)
        {
            return;
        }
        
        //用户处理握手后逻辑
        result = OnHandeShakeOver();
        if (!result)
        {
            return;
        }

        //处理实际数据帧部分
        while (1)
        {
            auto result2 = m_wsUnPack.UnPackFrame(this, m_nWaitMs);
            if (!result2)
            {
                break;
            }

            result = OnRecvText(result2.Get());
            if (!result)
            {
                break;
            }
        }
    }

protected:
    /**
    *@brief     握手过程
    */
    CPCResult<> HttpHandShake()
    {
        CPCResult<> result;
        CPCRestHttpUnPack  httpUnPack;
        httpUnPack.UnPackInit(true);
        while (1)
        {
            CPCBuffer package(PC_MAXSTACK);
            result = XRecv(m_nWaitMs, package);
            if (!result)
            {
                return result;
            }

            httpUnPack.UnPackUpdate(package.data(), package.length());
            if(!httpUnPack.m_LastResult)
            {
                //出错
                return result.SetFail("%s", httpUnPack.m_LastResult.ErrDesc());
            }
            else if (0 == httpUnPack.m_LastResult.Get())
            {
                //继续接收
                continue;
            }

            //完成
            auto wskey = httpUnPack.m_header.GetHeaderValue("Sec-WebSocket-Key");
            if (wskey.empty())
            {
                //出错，客户端没有传入密钥
                return result.SetFail("client req not find Sec-WebSocket-Key"); 
            }

            //协议
            std::string protocalline;
            auto protocalvalue = httpUnPack.m_header.GetHeaderValue("Sec-WebSocket-Protocol");
            if (!protocalvalue.empty())
            {
                protocalline = "Sec-WebSocket-Protocol: ";
                protocalline += protocalvalue;
                protocalline += "\r\n";
            }

            // 算出WEBSOCKET响应信息
            std::string acceptkey = CPCWebSocketUnPack::GetValidateCode(wskey);
            std::string response = "HTTP/1.1 101 Switching Protocols\r\n"
                "Upgrade: websocket\r\n"
                "Connection: Upgrade\r\n"
                "Sec-WebSocket-Accept: ";
            response = response + acceptkey + "\r\n" + protocalline + "\r\n";
            result = XSend(response.data(), response.length(), m_nWaitMs);
            if (!result)
            {
                return result;
            }

            break;
        }
        return result;
    }

protected:
    int m_nWaitMs;
    CPCWebSocketUnPack m_wsUnPack;
};




#endif    //___PCLIB_S_WEBSOCKET___